from sentence_transformers import SentenceTransformer

# 加载预训练模型（典型的Bi-Encoder）
model = SentenceTransformer("all-MiniLM-L6-v2")

# 编码单个句子
embedding1 = model.encode("How many people live in Berlin?")

# 编码多个句子
sentences = [
    "Berlin has 3.5 million inhabitants",
    "Paris has about 2.1 million residents",
]
embeddings = model.encode(sentences)

# 计算余弦相似度
from sklearn.metrics.pairwise import cosine_similarity

similarity = cosine_similarity([embedding1], embeddings)
print(similarity)
